import torch
from utils.overwatch import initialize_overwatch

logger = initialize_overwatch("vla_qat")


@torch.compiler.disable
def print_loss(x: torch.Tensor, y: torch.Tensor, name=""):
    with torch.inference_mode():
        loss = torch.nn.functional.mse_loss(x, y, reduction="none").detach().cpu()
        x = x.float()
        logger.info(
            f"{name[27:]}_loss:\n xmax: {x.max().item():.6f} xmean: {x.mean().item():.6f} x99: {torch.quantile(x, 0.99)} x90:{torch.quantile(x, 0.9)} \n xmax: {x.abs().max().item():.6f} xmean: {x.abs().mean().item():.6f} x99: {torch.quantile(x.abs(), 0.99)} x90: {torch.quantile(x.abs(), 0.9)} \n max: {loss.max().item():.6f}, mean: {loss.mean().item():.6f}"
        )
